extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, DeriveInput};

#[proc_macro_derive(DataTransferObject)]
pub fn data_transfer_object_derive(input: TokenStream) -> TokenStream {
    // 将输入的 token 解析为语法树
    let input = parse_macro_input!(input as DeriveInput);

    // 获取我们正在为其派生的结构体的名称
    let name = input.ident;

    // 构建 trait 实现的代码
    let expanded = quote! {
        impl DataTransferObject for #name {
            /// 将数据序列化为 JSON 字节数组
            fn to_json_bytes(&self) -> anyhow::Result<bytes::Bytes> {
                let json_string = serde_json::to_string(self)?;
                Ok(bytes::Bytes::from(json_string))
            }

            /// 从 JSON 字节数组反序列化数据
            fn from_json_bytes(bytes: &bytes::Bytes) -> anyhow::Result<Self> {
                let json_str = std::str::from_utf8(bytes)?;
                let data = serde_json::from_str(json_str)?;
                Ok(data)
            }

            /// 将数据序列化为二进制字节数组 (使用 bincode)
            fn to_binary_bytes(&self) -> anyhow::Result<bytes::Bytes> {
                let binary_data = bincode::serialize(self)?;
                Ok(bytes::Bytes::from(binary_data))
            }

            /// 从二进制字节数组反序列化数据 (使用 bincode)
            fn from_binary_bytes(bytes: &bytes::Bytes) -> anyhow::Result<Self> {
                let data = bincode::deserialize(bytes)?;
                Ok(data)
            }

            // TODO: 可以添加一个结构体指纹
        }
    };

    // 将生成的代码返回给编译器
    TokenStream::from(expanded)
}

/// 为 SandboxType 枚举自动实现 SandboxTrait 的宏
/// 此宏确保枚举的每个变体都必须实现 SandboxTrait
#[proc_macro_derive(SandboxTypeImpl)]
pub fn sandbox_type_impl_derive(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = input.ident;
    
    // 检查是否为枚举类型
    let variants = match input.data {
        syn::Data::Enum(data_enum) => data_enum.variants,
        _ => panic!("SandboxTypeImpl can only be used on enums"),
    };
    
    // 为每个枚举变体生成匹配分支
    let trait_methods = generate_sandbox_trait_methods(&name, &variants);
    
    // 生成编译时约束检查
    let constraint_checks = generate_constraint_checks(&variants);
    
    let expanded = quote! {
        // 编译时约束检查：确保所有变体类型都实现了 SandboxTrait
        #constraint_checks
        
        impl crate::core::sandbox::SandboxTrait for #name {
            #trait_methods
        }
    };
    
    TokenStream::from(expanded)
}

/// 生成 SandboxTrait 的方法实现
fn generate_sandbox_trait_methods(
    enum_name: &syn::Ident,
    variants: &syn::punctuated::Punctuated<syn::Variant, syn::token::Comma>,
) -> proc_macro2::TokenStream {
    // 为每个变体生成匹配分支
    let variant_names: Vec<_> = variants.iter().map(|v| &v.ident).collect();
    
    quote! {
        fn new(stage: eztrade_dto::stage::StageInfo) -> Self {
            // 根据 stage.stage_type 选择构建的沙盒类型
            match stage.stage_type.clone() {
                eztrade_dto::stage::StageType::SIM => {
                    #enum_name::SimpleSimulator(
                        crate::core::sandbox::simulate::SimpleSandboxSimulator::new(stage)
                    )
                }
                other => {
                    unimplemented!(
                        "Sandbox for stage_type {:?} is not implemented yet",
                        other
                    )
                }
            }
        }
        
        fn get_stage_phase(&self) -> eztrade_dto::stage::StagePhase {
            match self {
                #(#enum_name::#variant_names(inner) => inner.get_stage_phase(),)*
            }
        }
        
        fn get_stage_info(&self) -> eztrade_dto::stage::StageInfo {
            match self {
                #(#enum_name::#variant_names(inner) => inner.get_stage_info(),)*
            }
        }
        
        fn start(&self) -> anyhow::Result<()> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.start(),)*
            }
        }
        
        fn subscribe(&self, client_id: u64, subscribe: eztrade_dto::events::SubscribeEvent) -> anyhow::Result<()> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.subscribe(client_id, subscribe),)*
            }
        }
        
        fn unsubscribe(&self, client_id: u64, subscribe: eztrade_dto::events::SubscribeEvent) -> anyhow::Result<()> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.unsubscribe(client_id, subscribe),)*
            }
        }
        
        fn subscribe_global(&self, subscribe: eztrade_dto::events::GlobalEvent) -> anyhow::Result<()> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.subscribe_global(subscribe),)*
            }
        }
        
        fn unsubscribe_global(&self, subscribe: eztrade_dto::events::GlobalEvent) -> anyhow::Result<()> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.unsubscribe_global(subscribe),)*
            }
        }
        
        fn get_receiver(&self, client_id: u64) -> anyhow::Result<tokio::sync::broadcast::Receiver<eztrade_dto::Message>> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.get_receiver(client_id),)*
            }
        }
        
        fn process_message(&self, client_id: u64, message: eztrade_dto::Message) -> Option<eztrade_dto::Message> {
            match self {
                #(#enum_name::#variant_names(inner) => inner.process_message(client_id, message),)*
            }
        }
    }
}

/// 生成编译时约束检查
fn generate_constraint_checks(
    variants: &syn::punctuated::Punctuated<syn::Variant, syn::token::Comma>,
) -> proc_macro2::TokenStream {
    let checks = variants.iter().map(|variant| {
        // 获取变体的类型
        if let syn::Fields::Unnamed(fields) = &variant.fields {
            if let Some(field) = fields.unnamed.first() {
                let field_type = &field.ty;
                
                return quote! {
                    const _: fn() = || {
                        fn assert_sandbox_trait<T: crate::core::sandbox::SandboxTrait>() {}
                        assert_sandbox_trait::<#field_type>();
                    };
                };
            }
        }
        quote! {}
    });
    
    quote! {
        #(#checks)*
    }
}

/// 用于确保 Sandbox 实现必须满足 SandboxTrait 的约束宏
/// 在编译时检查类型是否实现了 SandboxTrait
#[proc_macro]
pub fn ensure_sandbox_trait(input: TokenStream) -> TokenStream {
    let type_name = parse_macro_input!(input as syn::Type);
    
    let expanded = quote! {
        const _: fn() = || {
            fn assert_impl_sandbox_trait<T: crate::core::sandbox::SandboxTrait>() {}
            assert_impl_sandbox_trait::<#type_name>();
        };
    };
    
    TokenStream::from(expanded)
}

/// 用于为新的 Sandbox 类型添加到 SandboxType 枚举的宏
/// 此宏会自动检查新类型是否实现了 SandboxTrait
#[proc_macro]
pub fn add_sandbox_variant(input: TokenStream) -> TokenStream {
    // 解析输入：variant_name(Type)
    let _input_str = input.to_string();
    
    // 这个宏主要用于文档和编译时提示
    // 实际的枚举扩展需要手动进行，但会有编译时检查
    let expanded = quote! {
        // 编译时检查：确保所有添加的类型都实现了 SandboxTrait
        compile_error!("请手动将新的 Sandbox 变体添加到 SandboxType 枚举中，并确保该类型实现了 SandboxTrait");
    };
    
    TokenStream::from(expanded)
}
